from pdbdataconnector import PDBDataConnector
from xml.etree.ElementTree import ElementTree

pdbxNS = "{http://pdbml.pdb.org/schema/pdbx-v40.xsd}"
class XMLDataConnector(PDBDataConnector) :
	def __init__(s, filepathORurl) :
		PDBDataConnector()
		s.tree = ElementTree()
		s.tree.parse( s.makeFileObject(filepathORurl) )
	def getCategory(s, cat, item=None, mode="normal") :
		'''overriding the parent method'''
		path = "%s%sCategory/%s%s"%(pdbxNS,cat,pdbxNS,cat)
		mult = 0 ; node = None
		for hit in s.tree.iterfind(path) :
			mult += 1
			if mult == 2 : break
			node = hit
		if mult == 0 : raise "No category found: " + cat
		if mode == "normal" :
			if mult == 1 : return XMLDataConnector.catRetMask(node, item)
			else : return XMLDataConnector.AnIterator(s.tree.iterfind(path), item)
		elif mode == "iter" :
			return XMLDataConnector.AnIterator(s.tree.iterfind(path), item)
		elif mode == "first" :
			return XMLDataConnector.catRetMask(node, item)
		else : raise NotImplementedError
	@staticmethod
	def catRetMask(node, item) :
		'''make map from node contents, or just text of child node with tag called item'''
		itemmap = XMLDataConnector.itemMap(node)
		if item == None : return itemmap
		return itemmap[item]
	@staticmethod
	def itemMap(node) :
		'''make a dictionary from tag name to tag text within this node'''
		retmap = {}
		for cn in node : retmap[cn.tag[len(pdbxNS):]] = cn.text
		for key,val in node.attrib.items() : retmap[key] = val
		return retmap
	class AnIterator :
		def __init__(s, catiter, item) : s.catiter, s.item = catiter, item
		def __iter__(s) : return s.forward()
		def forward(s) :
			for hit in s.catiter :
				yield XMLDataConnector.catRetMask(hit, s.item)
